from typing import Dict, List
from ..instances.ExampleFormatter import ExampleFormatter
from ..instances.MainTaskFormatter import MainTaskFormatter
from ..format_utils import *


class RETaskFormatter(MainTaskFormatter):
    """
    Task: Given sentence and entities, type the relation between them
    """

    def __init__(
        self,
        with_icl: bool = False,
        with_examples: bool = False,
        **kwargs,
    ) -> None:

        super().__init__(**kwargs)

        self.with_icl = with_icl
        self.with_examples = with_examples
        self.kwargs = kwargs

    def format_data_sample(
        self,
        data_item,
        instruction: str,
        options: List[str] = None,
    ) -> Dict:
        # sentence = data_item["sentence"]
        relations = data_item["relations"]

        context = self.format_context(data_item)

        instruction = add_instruction_prefix(instruction)
        options_joined = add_options_prefix(format_str_list_to_a_string(options))

        answer = self.format_answer(relations)
        # answer = super().format_answer(relations)

        examples_format = ""

        if self.with_icl:
            examples_format = ExampleFormatter(**self.kwargs).from_examples(
                examples=data_item["examples"]
            )

        if self.with_examples:
            example_relations = [
                r for exp in data_item["examples"] for r in exp["relations"]
            ]
            examples_format = add_examples_prefix(self.format_answer(example_relations))

        return {
            "instruction": instruction + " " + options_joined + " " + examples_format,
            "context": context + " " + Preffix.ANSWER.value,
            "labels": answer,
        }
